# -*- coding: UTF-8 -*-
"""
@项目名称：utiltools.py
@作   者：陆地起飞全靠浪
@创建日期：2025-09-18-15:55
"""
import os
import shutil
from openbabel import pybel
from pymol import cmd
from rdkit import Chem


def get_verify_smiles(df_valus_np):
    verify_smiles_list = []
    for line in df_valus_np:
        try:
            smiles = line[0]
            mol = Chem.MolFromSmiles(smiles)
            mol.GetNumAtoms()
            # 输出标准化后的 不带有手性SMILES
            Chem.MolToSmiles(mol, isomericSmiles=False, canonical=True, kekuleSmiles=True)
            verify_smiles_list.append(line)
        except:
            pass
    return verify_smiles_list


def ligand_format_conversion(input_file_path, output_file_path, input_type=None, output_type=None):
    '''
    :param input_file_path:
    :param output_file_path:
    :param input_type: pdb、mol2、sdf、cif
    :param output_type: pdb、mol2、sdf、cif
    :return:
    '''
    os.makedirs(os.path.dirname(output_file_path), exist_ok=True)
    if input_type not in ['pdb', 'mol2', 'sdf', 'cif']:
        input_type = os.path.splitext(input_file_path)[1][1:]
    if output_type not in ['pdb', 'mol2', 'sdf', 'cif']:
        output_type = os.path.splitext(output_file_path)[1][1:]
    if input_type in ['cif']:
        mol = next(pybel.readfile(input_type, input_file_path))
        output = pybel.Outputfile(output_type, output_file_path, overwrite=True)
        output.write(mol)
        output.close()
    else:
        cmd.delete('all')
        cmd.load(input_file_path, "all_")  # Warning: 'all' is a reserved keyword, appending underscore
        cmd.save(output_file_path, "all")
        cmd.delete('all')

    print(f'{input_type}转{output_type}完成，结果保存路径为：{output_file_path}')


if __name__ == '__main__':
    # for pdb_id in ['4bvw', '8tce', '8v8z', '8v9b']:
    #     input_file_path = f'/4T/AIDDserver/DataSrc/lpa_test/cif_src/{pdb_id}_s.cif'
    #     output_file_path = f'/4T/AIDDserver/DataSrc/lpa_test/{pdb_id}/{pdb_id}.pdb'
    #     ligand_format_conversion(input_file_path=input_file_path, output_file_path=output_file_path)
    #     shutil.copy(input_file_path, f'/4T/AIDDserver/DataSrc/lpa_test/{pdb_id}/{pdb_id}.cif')
    ligand_format_conversion(input_file_path='/4T/AIDDserver/DataSrc/测试用例/7eot.cif', output_file_path='/4T/AIDDserver/DataSrc/测试用例/7eot.pdb')
